# -*- coding: UTF-8 -*-
"""
@Project ：data_gateway
@File    ：lyyj_run_test.py
@IDE     ：PyCharm
@Author  ：zhangyiheng
@Date    ：2025/2/18 17:47
@Describe：多分类检测
"""
# ! python
# -*- coding: UTF-8 -*-
import time
import threading
from collections import deque
from datetime import datetime

import cv2
import numpy as np
import torch
from pathlib import Path
from ultralytics import YOLO
from celery_app.data_sms.tasks import send_alarm_task
import logging

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),  # 输出到控制台
    ]
)
logger = logging.getLogger(__name__)

# 模型配置
MODEL_PATH = "/workspace/models/lyyj_dac11m_safe_v4/weights/best.pt"
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO(MODEL_PATH).to(device)
logger.info(f"模型加载完成 | 设备: {device}")

# 分类配置
CLASS_CONFIG = {
    'names': {0: '持械', 1: '火患', 2: '打斗', 3: '人员', 4: '动物'},
    'colors': {0: (255, 0, 0), 1: (0, 0, 255), 2: (0, 255, 255), 3: (0, 255, 0), 4: (255, 255, 0)}
}

# 系统参数
CONFIG = {
    'fire_validation': {
        'enabled': True,  # 是否启用火患验证
        'gradient_threshold': 0.15,  # 亮度梯度阈值
        'flicker_frames': 5,  # 频域分析帧数
        'flicker_min_hz': 1,  # 最小闪烁频率(Hz)
        'flicker_max_hz': 10,  # 最大闪烁频率(Hz)
        'flicker_energy_ratio': 0.25,  # 频域能量占比阈值
        'high_conf_thresh': 0.7,  # 高置信度直接通过阈值
        'low_conf_thresh': 0.25,  # 低置信度直接过滤阈值
        'min_roi_size': 20  # 最小有效ROI尺寸
    },
    'frame_size': (640, 640),
    'base_conf': 0.30,  # 基础置信度（覆盖所有类别）
    'class_conf': {  # 分类置信度阈值
        0: 0.40,  # 持械
        1: 0.40,  # 火患
        2: 0.60,  # 打斗
        3: 0.50,  # 人员
        4: 0.50  # 动物
    },
    'alarm_interval': 60,  # 报警间隔(秒)
    'output_dir': Path("alarm_frames"),
    'min_detections': 10,  # 最小连续检测次数
    'detection_window': 10,  # 检测时间窗口(秒)
    'frame_skip': 2,  # 帧跳过间隔
    'reconnect_interval': 3  # 重连间隔(秒)
}

# 摄像头列表
CAMERAS = {
    "area1": "rtmp://192.168.12.20/live/34020000002000001031",
    "area2": "rtmp://192.168.12.20/live/34020000002000001017",
    "area3": "rtmp://192.168.12.20/live/34020000002000001016",
}


class CameraProcessor:
    def __init__(self, url, area_id):
        self.area_id = area_id
        self.url = url
        self.cap = None
        self.active = True
        self.last_alarms = {}  # 报警记录 {类名: {'times': [时间戳], 'last_sent': 时间}}
        self.frame_count = 0
        self.error_count = 0
        self.last_reconnect = 0
        self.connect()
        # self.last_frame = None  # 帧缓存

        self.last_frames = deque(maxlen=CONFIG['fire_validation']['flicker_frames'])
        self.fire_validator = FireValidator()  # 火焰检测

        # 创建输出目录
        CONFIG['output_dir'].mkdir(exist_ok=True)

        self._warmup_model()

        # 启动处理线程
        threading.Thread(target=self.process, daemon=True).start()

    def _warmup_model(self):
        """模型预热方法"""
        warmup_frame = np.zeros((*CONFIG['frame_size'], 3), dtype=np.uint8)  # 创建空白帧
        try:
            _ = model(warmup_frame,
                      imgsz=CONFIG['frame_size'],
                      device=device,
                      verbose=False)
            logger.info(f"[{self.area_id}] 模型预热完成")
        except Exception as e:
            logger.error(f"[{self.area_id}] 模型预热失败: {str(e)}")

    def connect(self):
        """建立摄像头连接"""
        try:
            self.cap = cv2.VideoCapture(self.url)
            if self.cap.isOpened():
                self.cap.set(cv2.CAP_PROP_FPS, 30)
                self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
                self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, CONFIG['frame_size'][0])
                self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CONFIG['frame_size'][1])
                logger.info(f"摄像头连接成功: {self.area_id}")
                return True
            logger.error(f"摄像头连接失败: {self.area_id}")
            return False
        except Exception as e:
            logger.error(f"连接异常: {self.area_id} - {str(e)}")
            return False

    def process(self):
        """主处理循环"""
        while self.active:
            try:
                # 1. 连接状态检查（增加延迟避免高频重连）
                if not self._check_camera_connection():
                    time.sleep(1)
                    continue

                # 2. 读取视频帧（增加错误计数重置）
                ret, frame = self.cap.read()
                if not ret:
                    self._handle_read_error()
                    continue
                self.error_count = 0  # 成功读取后重置错误计数

                # 维护帧缓存
                self.last_frames.append(frame.copy())

                # 帧跳过控制
                self.frame_count += 1
                if self.frame_count % CONFIG['frame_skip'] != 0:
                    continue

                # 4. 执行检测
                try:
                    results = self.detect(frame)
                    if results:  # 初步检测到目标
                        processed = self.post_process(results, frame)
                        # 只有处理后仍有有效目标时才触发报警检查
                        if processed:
                            self.check_alarms(processed, frame.copy())
                except torch.cuda.OutOfMemoryError:
                    self._handle_gpu_oom()
                    continue

            except Exception as e:
                logger.error(f"[{self.area_id}] 处理循环异常: {str(e)}", exc_info=True)
                time.sleep(1)

    def _check_camera_connection(self):
        """摄像头连接检查（解耦子方法）"""
        if self.cap and self.cap.isOpened():
            return True

        if time.time() - self.last_reconnect < CONFIG['reconnect_interval']:
            return False

        self.reconnect()
        return False

    def _handle_read_error(self):
        """处理视频读取错误"""
        self.error_count += 1
        logger.warning(f"[{self.area_id}] 视频读取失败 (累计错误: {self.error_count})")
        if self.error_count > 2:
            self.reconnect()

    def _handle_gpu_oom(self):
        """处理GPU内存不足"""
        logger.warning(f"[{self.area_id}] GPU内存不足，尝试清理...")
        torch.cuda.empty_cache()
        time.sleep(1)

    def reconnect(self):
        """重新连接摄像头"""
        if self.cap:
            self.cap.release()
        self.last_reconnect = time.time()
        self.error_count = 0
        self.connect()

    def detect(self, frame):
        """执行目标检测"""
        try:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = model(
                rgb_frame,
                imgsz=CONFIG['frame_size'],
                conf=CONFIG['base_conf'],
                device=device,
                iou=1.0,  # 禁用内置NMS
                verbose=False
            )
            return [{
                'xyxy': box.xyxy[0].cpu().numpy(),
                'conf': box.conf.item(),
                'cls_id': int(box.cls.item()),
                'time': time.time()
            } for box in results[0].boxes]
        except Exception as e:
            logger.error(f"检测失败: {self.area_id} - {str(e)}")
            return []

    def post_process(self, detections, frame):
        """检测结果后处理"""
        valid_boxes = []

        # light = self._get_light_condition(self.last_frame)

        # 连续三个检测都误报直接结束循环
        validate_weapon_count: int = 0
        validate_fire_count: int = 0
        validate_fight_count: int = 0
        for box in detections:
            # 保存原始检测信息
            raw_cls = box['cls_id']
            raw_conf = box['conf']

            cls_threshold = CONFIG['class_conf'].get(raw_cls, CONFIG['base_conf'])

            if raw_conf < cls_threshold:
                continue

            # 持械类特殊处理
            if raw_cls == 0:

                if self.validate_weapon(box):
                    logger.debug(f"持械误报过滤")

                    validate_weapon_count += 1
                    if validate_weapon_count >= 3:
                        return valid_boxes

                    continue

            elif raw_cls == 1:  # 火患类误报处理
                if not self.validate_fire(box, frame):

                    validate_fire_count += 1
                    if validate_fire_count >= 3:
                        return valid_boxes

                    continue

            elif raw_cls == 2:  # 打斗类别特殊处理
                if not self.validate_fight(box):
                    box['cls_id'] = 3  # 转为人员类别
                    raw_cls = box['cls_id']
                    # box['conf'] = raw_conf * 0.8  # 降低置信度

            # 应用分类置信度过滤
            if box['conf'] >= cls_threshold:
                valid_boxes.append(box)

        return valid_boxes

    def validate_fire(self, box, frame):
        """火患验证整合方法"""
        if not CONFIG['fire_validation']['enabled']:
            return True

        conf = box['conf']
        cfg = CONFIG['fire_validation']

        # 置信度分级处理
        if conf > cfg['high_conf_thresh']:
            return True
        if conf < cfg['low_conf_thresh']:
            return False

        # 提取当前ROI
        x1, y1, x2, y2 = map(int, box['xyxy'])
        current_roi = frame[y1:y2, x1:x2]

        # 检查ROI有效性
        if current_roi.size == 0 or \
                (x2 - x1) < cfg['min_roi_size'] or \
                (y2 - y1) < cfg['min_roi_size']:
            return False

        # 获取历史ROI序列
        roi_sequence = []
        for f in list(self.last_frames)[-cfg['flicker_frames']:]:
            roi_sequence.append(f[y1:y2, x1:x2])

        # 执行验证
        is_flickering = self.fire_validator.check_flicker(roi_sequence)
        has_gradient = self.fire_validator.check_gradient(current_roi)

        logger.debug(f"火患验证 | 闪烁:{is_flickering} 梯度:{has_gradient} 置信度:{conf:.2f}")
        # 必须满足：无车灯特征 并且 有火焰特征
        return (not is_flickering) and has_gradient

    def validate_weapon(self, box):
        x1, y1, x2, y2 = map(int, box['xyxy'])
        w, h = x2 - x1, y2 - y1

        # 条件1：排除固定坐标误报
        if (1400 <= x1 <= 1500) and (2100 <= x2 <= 2200) and (100 <= y1 <= 130) and (1270 <= y2 <= 1280):
            return True

        # （夜间提升至0.5）
        min_conf = 0.5 if self.is_night_time() else CONFIG['class_conf'].get(box['cls_id'], CONFIG['base_conf'])
        if box['conf'] < min_conf:
            logger.debug(
                f"持械低置信度过滤({'夜间' if self.is_night_time() else '日间'}): {box['conf']:.2f} < {min_conf}")
            return True

        # 条件2：通用尺寸限制
        return ((550 <= w <= 800) and (1000 <= h <= 1400))  # or (self.is_near_border(box['xyxy']))  # 边界检查

    def validate_fight(self, box):
        """验证打斗检测有效性"""
        x1, y1, x2, y2 = box['xyxy']
        w, h = x2 - x1, y2 - y1
        img_w, img_h = CONFIG['frame_size']

        # 基础条件检查
        conditions = [
            # box['conf'] > CONFIG['fighting_conf'],  # 置信度要求
            h > 0,  # 最小高度
            0.5 < w / h < 1.5,  # 宽高比范围
            not self.is_near_border(box['xyxy'])  # 边界检查
        ]

        return all(conditions)

    def is_night_time(self):
        """判断当前是否处于夜间时段（22:00-07:00）"""
        current_hour = datetime.now().hour
        return 18 <= current_hour or current_hour < 7

    def is_near_border(self, xyxy, thickness=30):
        """检查是否靠近边界"""
        x1, y1, x2, y2 = xyxy
        img_w, img_h = CONFIG['frame_size']
        return (x1 < thickness or (img_w - x2) < thickness or
                y1 < thickness or (img_h - y2) < thickness)

    def check_alarms(self, detections, frame):
        """检查报警条件
        参数:
            detections: 检测到的目标列表
            frame: 当前视频帧
        """
        current_time = time.time()
        alarms_to_send = {}

        logger.info(f"开始检查报警条件 | 区域: {self.area_id} | 检测到 {len(detections)} 个目标")

        # 先统计当前帧中出现的类别（去重）
        detected_classes = set()
        for box in detections:
            cls_id = box['cls_id']
            cls_name = CLASS_CONFIG['names'][cls_id]
            conf = box['conf']
            detected_classes.add(cls_name)

            logger.debug(
                f"检测目标 | 类别: {cls_name}({cls_id}) | "
                f"置信度: {conf:.2f} | 坐标: {box['xyxy']}"
            )

        # 对当前帧中出现的每个类别进行计数（每帧每类别只计一次）
        for cls_name in detected_classes:
            # 初始化或获取检测记录
            if cls_name not in self.last_alarms:
                self.last_alarms[cls_name] = {'times': [], 'last_sent': 0}
                logger.debug(f"初始化报警记录 | 类别: {cls_name}")

            records = self.last_alarms[cls_name]
            records['times'].append(current_time)

            # 清理过期记录
            valid_times = [t for t in records['times']
                           if current_time - t <= CONFIG['detection_window']]
            expired_count = len(records['times']) - len(valid_times)
            records['times'] = valid_times

            if expired_count > 0:
                logger.debug(
                    f"清理过期记录 | 类别: {cls_name} | "
                    f"清理数量: {expired_count} | 剩余有效: {len(valid_times)}"
                )

            # 检查报警条件
            detections_count = len(records['times'])
            time_since_last = current_time - records['last_sent']

            logger.debug(
                f"检查报警条件 | 类别: {cls_name} | "
                f"当前检测次数: {detections_count}/{CONFIG['min_detections']} | "
                f"距上次报警: {time_since_last:.1f}s/{CONFIG['alarm_interval']}s"
            )

            if (detections_count >= CONFIG['min_detections'] and
                    time_since_last >= CONFIG['alarm_interval']):

                # 获取该类别的最高置信度
                max_conf = max(b['conf'] for b in detections if CLASS_CONFIG['names'][b['cls_id']] == cls_name)

                logger.info(
                    f"满足报警条件 | 类别: {cls_name} | "
                    f"连续检测: {detections_count}次 | "
                    f"最高置信度: {max_conf:.2f}"
                )

                # 保存报警帧
                frame_path = self.save_alarm_frame(frame, detections)
                if not frame_path:
                    logger.error(f"报警帧保存失败 | 类别: {cls_name}")
                    continue

                # 记录要发送的报警
                alarms_to_send[cls_name] = {
                    'max_conf': max_conf,
                    'timestamp': current_time,
                    'frame_path': str(frame_path)
                }

                # 更新发送时间
                records['last_sent'] = current_time
                records['times'].clear()

                logger.info(
                    f"准备发送报警 | 类别: {cls_name} | "
                    f"最后报警时间更新: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(current_time))}"
                )

        # 发送报警
        if alarms_to_send:
            logger.info(
                f"开始发送报警 | 区域: {self.area_id} | "
                f"报警类别: {list(alarms_to_send.keys())} | "
                f"总报警数: {len(alarms_to_send)}"
            )
            self.send_alarms({
                'area_id': self.area_id,
                'alarms': alarms_to_send
            })
        else:
            logger.debug(f"无满足条件报警 | 区域: {self.area_id}")

    def save_alarm_frame(self, frame, detections):
        """保存报警帧"""
        try:
            # 绘制检测框
            for box in detections:
                x1, y1, x2, y2 = map(int, box['xyxy'])
                color = CLASS_CONFIG['colors'][box['cls_id']]
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

                # 添加标签
                label = f"{CLASS_CONFIG['names'][box['cls_id']]} {box['conf']:.2f}"
                cv2.putText(frame, label, (x1, y1 + 15),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

            # 保存图片
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            save_path = CONFIG['output_dir'] / f"{self.area_id}_{timestamp}.jpg"
            cv2.imwrite(str(save_path), frame)
            return save_path
        except Exception as e:
            logger.error(f"保存帧失败: {str(e)}")
            return None

    def send_alarms(self, alarms):
        """发送报警信息"""
        for cls_name, info in alarms.items():
            logger.info(
                f"报警触发 [{self.area_id}] {cls_name} "
                f"总数据: {info}"
            )

        try:
            # 转换数据格式
            formatted_data = {
                'area_id': alarms['area_id'],
                'alarms': [
                    {
                        'class': cls,
                        'confidence': info['max_conf'],
                        'timestamp': info['timestamp'],
                        'frame_path': info['frame_path']
                    }
                    for cls, info in alarms['alarms'].items()
                ]
            }

            send_alarm_task.delay(formatted_data)

            logger.info(
                f"[异步处理] 区域 {formatted_data.get('area_id')} 发送警报 =================================================== "
            )
        except Exception as e:
            logger.error(f"发送报警信息出错:{str(e)}")

    def release(self):
        """释放资源"""
        self.active = False
        if self.cap and self.cap.isOpened():
            self.cap.release()
        logger.info(f"摄像头关闭: {self.area_id}")


def cleanup(processors):
    """清理函数"""
    logger.info("开始系统清理...")
    for p in processors:
        p.release()
    if device == 'cuda':
        torch.cuda.empty_cache()
    logger.info("清理完成")


class FireValidator:
    def __init__(self):
        cv2.ocl.setUseOpenCL(True)

    def check_gradient(self, roi):
        """亮度梯度验证"""
        if roi.size == 0:
            return False

        try:
            gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY).astype(float) / 255
            grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
            grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
            gradient = np.sqrt(grad_x ** 2 + grad_y ** 2)
            return np.mean(gradient) > CONFIG['fire_validation']['gradient_threshold']
        except Exception as e:
            logger.error(f"梯度计算失败: {str(e)}")
            return False

    def check_flicker(self, roi_sequence):
        """频域闪烁分析"""
        if len(roi_sequence) < CONFIG['fire_validation']['flicker_frames']:
            return False

        brightness = []
        for roi in roi_sequence:
            if roi.size == 0:
                return False
            try:
                h, w = roi.shape[:2]
                center_roi = roi[max(0, h // 2 - 5):min(h, h // 2 + 5), max(0, w // 2 - 5):min(w, w // 2 + 5)]
                brightness.append(np.mean(cv2.cvtColor(center_roi, cv2.COLOR_BGR2HSV)[:, :, 2]))
            except Exception as e:
                logger.warning(f"ROI处理异常: {str(e)}")
                continue

        if len(brightness) < 3:  # 至少需要3个有效亮度值
            return False

        try:
            fft = np.fft.fft(brightness)
            freqs = np.fft.fftfreq(len(brightness), d=1 / 30)
            mask = (np.abs(freqs) >= CONFIG['fire_validation']['flicker_min_hz']) & \
                   (np.abs(freqs) <= CONFIG['fire_validation']['flicker_max_hz'])
            energy_ratio = np.sum(np.abs(fft[mask])) / (np.sum(np.abs(fft)) + 1e-6)  # 避免除零
            return energy_ratio > CONFIG['fire_validation']['flicker_energy_ratio']
        except Exception as e:
            logger.error(f"FFT分析失败: {str(e)}")
            return False


if __name__ == '__main__':
    processors = []
    try:
        # 初始化摄像头处理器
        for area_id, url in CAMERAS.items():
            processor = CameraProcessor(url, area_id)
            processors.append(processor)

        # 主循环
        while True:
            time.sleep(1)

    except KeyboardInterrupt:
        logger.info("接收到终止信号")
    except Exception as e:
        logger.critical(f"主程序异常: {str(e)}", exc_info=True)
    finally:
        cleanup(processors)
        logger.info("系统安全退出")
